/*
maths.c is part of 
ngstat (c) 2012-2013 Emanuele Raineri
emanuele.raineri@gmail.com


This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.


this program implements the formulae contained in
Luca Ferretti et al: 
Neutrality tests from high-throughput resequencing...
and
Luca Ferretti, Emanuele Raineri, Sebastian RO:

Neutrality tests for sequences with missing data

**************************************************/

#include "maths.h"

int min3(int a, int b, int c){
  return min(min(a,b),c);
}

int max3(int a, int b, int c){
  return max(max(a,b),c);
}


float lnfact(int n){
  static float table[1000];
  return (table[n]>0)?table[n]:(table[n]=lgamma(n+1)); 
}

float lnbinomial(int n, int k){
  static float table[100*100];
  if (n<0 || k<0)  {return -INFINITY;}
  if (k>n) {return -INFINITY;}
  //return lnfact(n)-lnfact(k)-lnfact(n-k);
  return (table[n*100+k]>0)?table[n*100+k]:(table[n*100+k]=lnfact(n)-lnfact(k)-lnfact(n-k)); 
}

float lnmultinomial(int n, int k1, int k2){
  if (n-k1-k2<0  ) {return -INFINITY;}
  if (n <0 || k1<0 || k2<0) {return -INFINITY;}
  return lnfact(n)-lnfact(k1)-lnfact(k2)-lnfact(n-k1-k2);
}

float an(int n){
  // sum of 1/i
  if (0==n) {fprintf(stderr,"a:invalid n\n");exit(1);}
  int i;
  float s=0.0;
  for (i=1;i<n;i++){
    s+=1.0/i;
  }
  return s;
}

float beta(int n,int i){
   int d1=n-i;
   if (d1==0) return 0;
   return 2.0*n*(an(n+1)-an(i))/(d1*(d1+1))-2.0/d1; 
}

float sigma(int i, int j, int n){
    /*
    i,j mutation sizes
    n number of sequences
    */
  int tmp;
  if (i<j) {tmp=i;i=j;j=tmp;}
  if (i==j){ 
    if ( 2*i < n) return beta(n,i+1);
    if ( 2*i == n ) return 2.0*(an(n)-an(i))/(n-i) - 1.0/(i*i);
    return beta(n,i)-1.0/(i*i);
  }
  if (i+j<n) return (beta(n,i+1)-beta(n,i))/2.0;
  if (i+j==n) return  ((an(n)-an(i))/(n-i))+((an(n)-an(j))/(n-j))-(beta(n,i)+beta(n,j+1))/2.0-1.0/(i*j);
  return (beta(n,j)-beta(n,j+1))/2.0-1.0/(i*j);
}

float omegad(int i,int x,int* n,int l){
    float sum=0.0;
    int y;
    for(y=0;y<l;y++){
        sum+=hn(n[y]-1);
    }
    float term1=2.0*(n[x]-i)/(n[x]*(n[x]-1.0));
    float term2=-1.0/(((float)i/l)*sum);
    return term1+term2;
}

float hn( int n ){
    //computes (rather naively) the
    //n-th harmonic number
    int i;
    float s =0.0;
    for(i=1;i<=n;i++) s+=1.0/i;
    return s; 	
}

float ngscov( int* n, int** r, float theta, int i, int j, int x, int y ) {
    int cx= round(sumc(r,x,n));
    int cy= round(sumc(r,y,n));
    if (cx<cy) { int tmp=cx;cx=cy;cy=tmp; }
    int l; float sum=0.0;
    for (l=j;l<=cx-cy+j;l++){
        sum+=sigma(i,l,cx)*exp(lnbinomial(l,i)+lnbinomial(cx-l,cy-j));
    }
    return theta*theta*sum/(exp(lnbinomial(cx,cy)));
}

float p_product(int x, int nind,  int *v, int** r){
    /* x is the position */
    float p = 1.0,tmp;
    int k;
    for ( k = 0 ; k < nind ; k++ ){
        if (v[k]==1) {tmp=pow(2,-r[x][k]+1 );}
        else {tmp = 1.0 - pow(2,-r[x][k]+1 );}
        p*=tmp;
    }
    return p;
}

float pc_j(int j, int x, int nind,  int** r){
    /*
    x is the position
    v[i] n of sequenced alleles for individual i      
    */
    int index=nind-1, *v,i,tmpsum=0;
    float pc = 0.0;
    v=malloc(nind*sizeof(int));
    for(i=0; i<nind; i++) v[i]=1;
    
    tmpsum = nind;
    if ( tmpsum == j ) pc+=p_product(x,nind,v,r);

    while( index>=0 ){
        if ( v[index] == 1 ){

            v[index]++;
            tmpsum=0;
            for(i=0;i<nind;i++) {
                /* fprintf(stderr,"=>v[%d](j=%d)=%d ",i,j,v[i]); */
                tmpsum+=v[i];
            
            }
            
            /* fprintf(stderr,"\n"); */
            index=nind-1;
            if (tmpsum!=j) continue; 
            else pc += p_product(x,nind,v,r);
        
        } else {
            
            v[index]=1;
            index--;
        
        }
    }
    return pc;
}

int n_segreg_sites(int** a, int* n, int ws){
    int x,k,s=0;
    for (x=0; x<ws; x++){
        for (k=0; k< n[x]; k++){
            if (a[x][k]==1) {
                s++;
                break;
            }
        }
    }
    return s;
}

float theta_w(int** a, int* n , int ws,  int** r, int* s){
    *s = n_segreg_sites( a, n,  ws ) ; 
    int x,j;
    float sum_denom=0.0;
    for(x=0;x<ws;x++){
        for(j=2; j<= 2*n[x]; j++){
            sum_denom += pc_j(j, x, n[x] , r)*an(j);
        }
    }
    return *s/sum_denom;
}

float est_seq_chr(int r){
	return 2*(1-pow(2.0,-r));
}

float sumc(int** r, int x, int* n){
    // n num individuals
    //r read depth
    float sum=0.0;
    int k, nind=n[x];
    for(k=0;k<nind;k++){
        sum+=est_seq_chr(r[x][k]);
    }
    return sum;
}


int est_gtype(int m,int r){
	/* inferred allele count 
        r read depth
	m number of minor alleles 
        */
	if (0==m) return 0;
	if (r==m) return 2;
	return 1;
}

float fhat( int x,  int** a, float** c, int nind, float* var ){
	/*note : only one for loop to build both sums
	(minor) allele frequency estimation
	it also returns the variance */
	int k;
	float sumc=0.0,sum1=0.0,sum2=0.0,tmp=0.0;
	float f;
	for( k=0; k<nind; k++ ){ /* loop over the individuals */
	
                tmp=1.0/(3.0-c[x][k]);
		sum1+=tmp*a[x][k];
		sum2+=tmp;
		sumc+=c[x][k];	
	}

	f = sum1/(2.0*sum2);
	*var = f*(1.0-f)*sumc/(2.0*tmp*tmp);
	return f;
}

float variance_tajima( int l, int nind, int* n, int** r, float theta ){

    int i,j,x,y;
    float term1=0.0,term2=0.0;
    float alphai,alphaj,cov;
    
    for(x=0;x<l;x++){
        for(i=1;i<=2*nind-1;i++){
            alphai=omegad(i,x,n,l);
            term1+=(1.0/i)*(alphai*alphai);
            for(y=0;y<l;y++){
                if (x==y) continue;
                for(j=1;j<=2*nind-1;j++){
                    alphaj = omegad( j, y , n, l );
                    cov = ngscov( n , r , theta, i , j, x, y );
                    term2 += i*j*alphai*alphaj*cov;
                }
            }
        }
    }

    return theta*term1 + term2;

}

float theta_pi(int** r,int **a,float **c, int* n, int ws){
        /*
        a number of alt alleles 0,1,2 for a given position and individual
        c number of chromosomes sequenced for a given position and individual
        n total number of individuals at position x
        r read depth for given position and individual
        returns tajima's pi i.e. avg number of nucleotide differences
        between  pairs of sequences.
        */
	int k,l,x,ne;

	float sum1,sum2,sum3,sum4=0.0,sumc;
        
        for (x=0;x<ws;x++){
            
            sum1=0.0; sum2=0.0; sum3=0.0; sumc=0.0; ne=0;

            for(k=0; k<n[x]; k++ ){

                if ( r[x][k] == 0 ) continue; else ne++;

		sum1 += a[x][k]*(2-a[x][k]);
		
                for (l=0; l<n[x]; l++){
                    
                    if (r[x][l]==0 || k==l) continue;
		    
                    sum2 += c[x][k]*c[x][l]*a[x][k]*(2.0-a[x][l])/4.0;

		    sum3 += c[x][k]*c[x][l]/2.0;
		
                }
	        sumc += c[x][k];
            
            }
            
            sum4 += ( sum1 + sum2 ) / ( sumc - ne + sum3 ); 
        }
	return ( 1.0 / ws ) * sum4;
}
